import torch
from torchtext.vocab import build_vocab_from_iterator
from torch.nn.utils.rnn import pad_sequence

def build_vocab(dataset):
    '''
    构建词汇表
    :param dataset:
    :return:
    '''
    specials = ["<unk>", "<pad>"]
    # (tokens, label)
    ite_text = map(lambda x:x[0], dataset)

    #构建词汇表
    vocab_text = build_vocab_from_iterator(ite_text,  min_freq=2, specials=specials)

    #未记载词
    vocab_text.set_default_index(vocab_text["<unk>"])
    return vocab_text

def collate_batch(batch, vocab):
    token_list = []
    labels = []

    for text, label in batch:
        text_idx = [vocab[token] for token in text]
        text_ts = torch.tensor(text_idx, dtype=torch.long)
        token_list.append(text_ts)

        label_ts = torch.tensor(label - 1, dtype=torch.long)
        labels.append(label_ts)

    pad_idx = vocab['<pad>']
    padded_text = pad_sequence(token_list, True, pad_idx)
    labels = torch.stack(labels)

    return padded_text, labels



